package com.gallop.mq.core.template;

import com.gallop.mq.MQProperties;
import com.gallop.mq.MQTemplate;
import com.gallop.mq.core.MQExecutorProvider;
import com.gallop.mq.core.MQMessageCodec;
import com.gallop.mq.core.MQMessageDispatcher;
import com.gallop.mq.core.dispatcher.DispatchContext;
import com.gallop.mq.core.exception.MQException;
import com.gallop.mq.core.template.handler.DeliveryContext;
import com.gallop.mq.core.template.handler.DeliveryType;
import com.gallop.mq.core.template.handler.MessageDeliveryHandler;
import com.gallop.mq.core.template.handler.registry.DynamicProxyListenerCreator;
import com.gallop.mq.core.template.handler.registry.MessageDeliveryHandlerRegistry;
import com.gallop.mq.core.template.handler.registry.RegistryHolder;
import com.gallop.mq.core.template.handler.registry.RegistryPreprocessor;
import com.google.common.collect.Lists;
import com.google.common.collect.Maps;
import lombok.AllArgsConstructor;
import lombok.Data;
import lombok.Getter;
import lombok.Setter;
import lombok.extern.slf4j.Slf4j;
import org.apache.kafka.clients.consumer.ConsumerConfig;
import org.apache.kafka.clients.consumer.ConsumerRecord;
import org.apache.kafka.clients.consumer.ConsumerRecords;
import org.apache.kafka.clients.consumer.KafkaConsumer;
import org.apache.kafka.clients.producer.KafkaProducer;
import org.apache.kafka.clients.producer.ProducerConfig;
import org.apache.kafka.clients.producer.ProducerRecord;
import org.apache.kafka.common.header.Header;
import org.apache.kafka.common.header.Headers;
import org.apache.kafka.common.header.internals.RecordHeader;
import org.apache.kafka.common.header.internals.RecordHeaders;
import org.apache.kafka.common.serialization.ByteArrayDeserializer;
import org.apache.kafka.common.serialization.ByteArraySerializer;
import org.apache.kafka.common.serialization.StringDeserializer;
import org.apache.kafka.common.serialization.StringSerializer;
import org.springframework.util.StringUtils;

import java.nio.charset.StandardCharsets;
import java.time.Duration;
import java.util.*;
import java.util.concurrent.*;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicInteger;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

@Getter
@Setter
@Slf4j
public class KafkaMQTemplate implements MQTemplate, RegistryHolder {

    public static final String HEADER_KEY_CORRELATION_ID = "correlationId";
    public static final String HEADER_KEY_REPLY_TO = "replyTo";
    public static final String HEADER_KEY_MSG_TYPE = "type";
    private MQMessageCodec codec;
    private MQMessageDispatcher dispatcher;
    private MQProperties properties;
    private MQExecutorProvider executorProvider;

    private String localReplyQueue;

    private Map<String, PendingReply> pendingReplyMap = Maps.newConcurrentMap();
    private KafkaProducer<String, byte[]> producer;

    private KafkaListenerContext listenerContext = new KafkaListenerContext();

    private AtomicBoolean working = new AtomicBoolean();

    private MessageDeliveryHandlerRegistry handlerRegistry;

    @Override
    public void init(MQProperties properties, MQMessageDispatcher dispatcher, MQMessageCodec codec, MQExecutorProvider executorProvider, MessageDeliveryHandlerRegistry handlerRegistry) throws MQException {
        setProperties(properties);
        setCodec(codec);
        setDispatcher(dispatcher);
        setExecutorProvider(executorProvider);
        listenerContext.start();
        this.initRegistry(handlerRegistry);


        // local producer
        Properties producerProperties = new Properties();
        producerProperties.put(ProducerConfig.BOOTSTRAP_SERVERS_CONFIG, properties.getHosts());
        producerProperties.put(ProducerConfig.KEY_SERIALIZER_CLASS_CONFIG, StringSerializer.class);
        producerProperties.put(ProducerConfig.VALUE_SERIALIZER_CLASS_CONFIG, ByteArraySerializer.class);

        producer = new KafkaProducer<>(producerProperties);
        // 初始化rpc回调监听
        localReplyQueue = "rpc-reply-" + UUID.randomUUID().toString().replaceAll("-", "");
        Properties consumerProperties = new Properties();
        consumerProperties.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, properties.getHosts());
        consumerProperties.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
        consumerProperties.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class);
        consumerProperties.put(ConsumerConfig.GROUP_ID_CONFIG, localReplyQueue);
        KafkaConsumer<String, byte[]> localReplyConsumer = new KafkaConsumer<>(consumerProperties);
        localReplyConsumer.subscribe(Collections.singletonList(localReplyQueue));

        // 收到rpc响应
        listenerContext.register(localReplyQueue, localReplyConsumer,
                handlerRegistry.createListener(DeliveryContext.builder().mqTemplate(this).build()));

        // 初始化扫描topic监听
        dispatcher.getAllContainers().forEach(container -> {
            Properties p = new Properties();
            p.put(ConsumerConfig.BOOTSTRAP_SERVERS_CONFIG, properties.getHosts());
            p.put(ConsumerConfig.KEY_DESERIALIZER_CLASS_CONFIG, StringDeserializer.class);
            p.put(ConsumerConfig.VALUE_DESERIALIZER_CLASS_CONFIG, ByteArrayDeserializer.class);
            p.put(ConsumerConfig.GROUP_ID_CONFIG, StringUtils.hasText(container.getGroup()) ? container.getGroup() : properties.getKafka().getDefaultTopicGroup());
            KafkaConsumer<String, byte[]> consumer = new KafkaConsumer<>(p);
            consumer.subscribe(Collections.singletonList(container.getTopic()));

            String id = container.getTopic() + container.getGroup();
            listenerContext.register(id, consumer, handlerRegistry.createListener(DeliveryContext.builder()
                    .mqTemplate(this)
                    .container(container)
                    .build()));
        });

        working.set(true);
        log.info("KafkaMQTemplate init");
    }

    private String parseHeaderValue(Headers headers, String key) {
        Header header = headers.lastHeader(key);
        if (header == null) {
            return null;
        }
        return new String(header.value(), StandardCharsets.UTF_8);
    }

    @Override
    public void onDispose() {
        if (this.producer != null) {
            this.producer.close();
        }
        listenerContext.close();
    }

    @Override
    public void send(Object message, String topic) {
        if (!working.get()) {
            log.error("MQTemplate is not working");
            return;
        }
        ProducerRecord<String, byte[]> record = createProducerRecord(topic, message);
        this.producer.send(record);
    }

    @Override
    public <P> void request(Object message, String topic, Duration requestTimeout, Consumer<P> onResponse) {
        request(message, topic, requestTimeout, onResponse, null);
    }

    @Override
    @SuppressWarnings("unchecked")
    public <P> void request(Object message, String topic, Duration requestTimeout, Consumer<P> onResponse, BiConsumer<Object, Exception> onError) {
        if (!working.get()) {
            log.error("MQTemplate is not working");
            return;
        }
        String correlationId = "rpc-correlation-" + UUID.randomUUID().toString().replaceAll("-", "");
        PendingReply pendingReply = new PendingReply();
        pendingReply.setSavedReplyTo(localReplyQueue);
        pendingReply.setSavedCorrelation(correlationId);
        try {
            this.pendingReplyMap.put(correlationId, pendingReply);
            Object response = doRequest(message, topic, pendingReply, requestTimeout);
            onResponse.accept((P) response);
        } catch (Exception e) {
            log.error("", e);
            this.pendingReplyMap.remove(correlationId);
            pendingReply.returned(e);
            if (onError != null) {
                onError.accept(message, e);
            }
        }
    }

    @Override
    @SuppressWarnings("unchecked")
    public <P> P request(Object message, String topic, Duration receiveTimeout) throws InterruptedException {
        if (!working.get()) {
            log.error("MQTemplate is not working");
            return null;
        }
        String correlationId = "rpc-correlation-" + UUID.randomUUID().toString().replaceAll("-", "");
        PendingReply pendingReply = new PendingReply();
        pendingReply.setSavedReplyTo(localReplyQueue);
        pendingReply.setSavedCorrelation(correlationId);
        try {
            this.pendingReplyMap.put(correlationId, pendingReply);
            Object response = doRequest(message, topic, pendingReply, receiveTimeout);
            return (P) response;
        } catch (Exception e) {
            log.error("", e);
            this.pendingReplyMap.remove(correlationId);
            pendingReply.returned(e);
        }
        return null;
    }

    @SuppressWarnings("unchecked")
    private <P> P doRequest(Object message, String topic, PendingReply pendingReply, Duration receiveTimeout) {
        ProducerRecord<String, byte[]> record = createProducerRecord(topic, message, pendingReply.getSavedReplyTo(), pendingReply.getSavedCorrelation(), DeliveryType.REQUEST);
        this.producer.send(record);
        Object response;

        if (receiveTimeout == null && properties.getRequestTimeout() > 0) {
            receiveTimeout = Duration.ofMillis(properties.getRequestTimeout());
        }
        if (receiveTimeout == null) {
            response = pendingReply.get();
        } else {
            response = pendingReply.get(receiveTimeout.toMillis(), TimeUnit.MILLISECONDS);
        }
        return (P) response;
    }


    private ProducerRecord<String, byte[]> createProducerRecord(String topic, Object message) {
        return createProducerRecord(topic, message, null, null, DeliveryType.SEND);
    }

    private ProducerRecord<String, byte[]> createProducerRecord(String topic,
                                                                Object message,
                                                                String replyTo,
                                                                String correlation, DeliveryType deliveryType) {
        List<Header> headers = Lists.newLinkedList();
        headers.add(new RecordHeader(HEADER_KEY_MSG_TYPE, deliveryType.name().getBytes(StandardCharsets.UTF_8)));
        if (StringUtils.hasLength(replyTo) && StringUtils.hasLength(correlation)) {
            headers.add(new RecordHeader(HEADER_KEY_REPLY_TO, replyTo.getBytes(StandardCharsets.UTF_8)));
            headers.add(new RecordHeader(HEADER_KEY_CORRELATION_ID, correlation.getBytes(StandardCharsets.UTF_8)));
        }
        return new ProducerRecord<>(topic, null, (String) null, this.codec.encode(message), new RecordHeaders(headers));
    }

    @Override
    public MessageDeliveryHandlerRegistry.InitializerBuilder registryInitializer() {
        return MessageDeliveryHandlerRegistry.InitializerBuilder.newBuilder()
                .setPreprocessor(RegistryPreprocessor.<Consumer<ConsumerRecord<String, byte[]>>, ConsumerRecord<String, byte[]>>builder()
                        .listenerCreator(DynamicProxyListenerCreator.ofGenericInterface(handlerRegistry,
                                Consumer.class))
//                .listenerCreator(context -> record -> {
//                    handlerRegistry.handleDelivery(context, record);
//                })
                        .deliveryTypeProvider(args -> {
                            ConsumerRecord<String, byte[]> record = args.getArgAs(0);
                            return DeliveryType.valueOf(parseHeaderValue(record.headers(), HEADER_KEY_MSG_TYPE));
                        })
                        .rawMessageProvider(args -> args.getArgAs(0))
                        .mqMessageProvider(record -> this.codec.decode(record.value()))
                        .rpcInfoProvider(args -> {
                            ConsumerRecord<String, byte[]> record = args.getArgAs(0);
                            String correlationId = parseHeaderValue(record.headers(), HEADER_KEY_CORRELATION_ID);
                            String replyTo = parseHeaderValue(record.headers(), HEADER_KEY_REPLY_TO);
                            return MessageDeliveryHandler.StringRpcInfo.builder()
                                    .correlationId(correlationId)
                                    .replyTo(replyTo)
                                    .build();
                        }).build())
                .registerHandler(DeliveryType.SEND, (deliveryContext, deliveryType, rpcInfo, incoming) ->
                        dispatcher.dispatch(DispatchContext.builder()
                                .message(incoming)
                                .topic(deliveryContext.getContainer().getTopic())
                                .group(deliveryContext.getContainer().getGroup())
                                .build()))
                .registerHandler(DeliveryType.REQUEST, (deliveryContext, deliveryType, rpcInfo, incoming) -> {
//                    MessageDeliveryHandler.StringRpcInfo stringRpcInfo = rpcInfo.as();
                    dispatcher.dispatch(DispatchContext.builder()
                            .message(incoming)
                            .topic(deliveryContext.getContainer().getTopic())
                            .group(deliveryContext.getContainer().getGroup())
                            .replyTo((String) rpcInfo.getReplyTo())
                            .correlationId(rpcInfo.getCorrelationId())
                            .rpcResponse((correlationId, replyTo, response) -> {
                                // rpc response
                                ProducerRecord<String, byte[]> respRecord = createProducerRecord(replyTo,
                                        response,
                                        replyTo,
                                        correlationId, DeliveryType.RESPONSE);
                                producer.send(respRecord);
                            }).build());
                })
                .registerHandler(DeliveryType.RESPONSE, (deliveryContext, deliveryType, rpcInfo, incoming) -> {
                    PendingReply pendingReply = this.pendingReplyMap.remove(rpcInfo.getCorrelationId());
                    if (pendingReply == null) {
                        log.error("rpc reply not found. record={},replyTo={},correlationId={},message={}",
                                rpcInfo.getReplyTo(),
                                localReplyQueue,
                                rpcInfo.getCorrelationId(),
                                incoming);
                        return;
                    }
                    pendingReply.reply(incoming);
                });
    }

    /**
     * 本地Consumer监听消息拉取所需上下文
     */
    @Data
    @AllArgsConstructor
    private static class KafkaListener {

        private AtomicBoolean pollingState;
        private KafkaConsumer<String, byte[]> consumer;
        private String id;
        private Consumer<ConsumerRecord<String, byte[]>> consumeFunction;
    }

    /**
     * 本地Consumer监听管理
     */
    @Data
    private class KafkaListenerContext {
        private Map<String, KafkaListener> listeners = Maps.newConcurrentMap();
        private ScheduledExecutorService scheduledExecutorService = new ScheduledThreadPoolExecutor(1, new KafkaListenerThreadPoolFactory());
        private ExecutorService workers;

        public void register(String id, KafkaConsumer<String, byte[]> consumer, Consumer<ConsumerRecord<String, byte[]>> consumeFunction) {
            KafkaListener listener = new KafkaListener(new AtomicBoolean(), consumer, id, consumeFunction);
            this.listeners.put(listener.getId(), listener);
        }

        public void start() {
            workers = executorProvider.create();
            // 启动定时任务拉取监听消息
            this.scheduledExecutorService.scheduleAtFixedRate(() ->
                    listeners.forEach((id, ctx) -> asyncPoll(ctx)), 0L, properties.getKafka().getMsgPollInterval(), TimeUnit.MILLISECONDS);
        }

        public void close() {
            scheduledExecutorService.shutdown();
            workers.shutdown();

            listeners.forEach((id, ctx) -> ctx.getConsumer().close());
            listeners.clear();
        }

        private void asyncPoll(KafkaListener ctx) {
            if (ctx.getPollingState().get()) {
                return;
            }
            workers.execute(() -> {
                if (!ctx.getPollingState().compareAndSet(false, true)) {
                    return;
                }
                try {
                    ConsumerRecords<String, byte[]> records = ctx.getConsumer().poll(Duration.ofMillis(properties.getKafka().getMsgPollTimeout()));
                    if (records == null) {
                        return;
                    }
                    if (log.isTraceEnabled()) {
                        log.trace("records poll:{},{}", records.count(), listeners.size());
                    }
                    records.forEach(r -> {
                        try {
                            ctx.getConsumeFunction().accept(r);
                        } catch (Exception e) {
                            log.error("", e);
                        }
                    });
                } finally {
                    ctx.getPollingState().set(false);
                }
            });
        }
    }

    private static class KafkaListenerThreadPoolFactory implements ThreadFactory {

        private static final String threadName = "kafka-listener-client-pool";
        private final ThreadGroup threadGroup;
        private final AtomicInteger threadCount = new AtomicInteger();


        public KafkaListenerThreadPoolFactory() {
            this.threadGroup = new ThreadGroup(threadName);
            threadGroup.setDaemon(false);
            threadGroup.setMaxPriority(Thread.NORM_PRIORITY);
        }

        @Override
        public Thread newThread(Runnable r) {
            return new Thread(threadGroup, r, threadName + "-worker-" + threadCount.getAndIncrement());
        }
    }
}
